import torch 
segrap_subset = {
    'Brain': [1, 2, 3, 4, 5, 6, 7, 8, 9],
    "BrainStem": 2,
    "Chiasm": 3,
    "TemporalLobe_L": [4, 6],
    "TemporalLobe_R": [5, 7],
    "Hippocampus_L": [8, 6],
    "Hippocampus_R": [9, 7],
    'Eye_L': [10, 12],
    'Eye_R': [11, 13],
    "Lens_L": 12,
    "Lens_R": 13,
    "OpticNerve_L": 14,
    "OpticNerve_R": 15,
    "MiddleEar_L": [18, 16, 20, 24, 28, 30],
    "MiddleEar_R": [19, 17, 21, 25, 29, 31],
    "IAC_L": 18,
    "IAC_R": 19,
    "TympanicCavity_L": [22, 20],
    "TympanicCavity_R": [23, 21],
    "VestibulSemi_L": [26, 24],
    "VestibulSemi_R": [27, 25],
    "Cochlea_L": 28,
    "Cochlea_R": 29,
    "ETbone_L": [32, 30],
    "ETbone_R": [33, 31],
    "Pituitary": 34,
    "OralCavity": 35,
    "Mandible_L": 36,
    "Mandible_R": 37,
    "Submandibular_L": 38,
    "Submandibular_R": 39,
    "Parotid_L": 40,
    "Parotid_R": 41,
    "Mastoid_L": 42,
    "Mastoid_R": 43,
    "TMjoint_L": 44,
    "TMjoint_R": 45,
    "SpinalCord": 46,
    "Esophagus": 47,
    "Larynx": [48, 49, 50, 51],
    "Larynx_Glottic": 49,
    "Larynx_Supraglot": 50,
    "PharynxConst": [51, 52],
    "Thyroid": 53,
    "Trachea": 54}


def convert_labels(labels):
    result = [(labels == 1) | (labels == 2) | (labels == 3) | (labels == 4) | (labels == 5) | (labels == 6) | (labels == 7) | (labels == 8) | (labels == 9), 
                (labels == 2),
                (labels == 3),
                (labels == 4) | (labels == 6),
                (labels == 5) | (labels == 7),
                (labels == 8) | (labels == 6),
                (labels == 9) | (labels == 7),
                (labels == 10) | (labels == 12),
                (labels == 11) | (labels == 13),
                (labels == 12),
                (labels == 13),
                (labels == 14),
                (labels == 15),
                (labels == 18) | (labels == 16) | (labels == 20) | (labels == 24) | (labels == 28) | (labels == 30),
                (labels == 19) | (labels == 17) | (labels == 21) | (labels == 25) | (labels == 29) | (labels == 31),
                (labels == 18),
                (labels == 19),
                (labels == 22) | (labels == 20),
                (labels == 23) | (labels == 21),
                (labels == 26) | (labels == 24),
                (labels == 27) | (labels == 25),
                (labels == 28),
                (labels == 29),
                (labels == 32) | (labels == 30),
                (labels == 33) | (labels == 31),
                (labels == 34),
                (labels == 35),
                (labels == 36),
                (labels == 37),
                (labels == 38),
                (labels == 39),
                (labels == 40),
                (labels == 41),
                (labels == 42),
                (labels == 43),
                (labels == 44),
                (labels == 45),
                (labels == 46),
                (labels == 47),
                (labels == 48) | (labels == 49) | (labels == 50) | (labels == 51),
                (labels == 49),
                (labels == 50),
                (labels == 51) | (labels == 52),
                (labels == 53),
                (labels == 54),
                ]

    return torch.cat(result, dim=1).float()



from models.bit_diffusion import decimal_to_bits, bits_to_decimal
if __name__ == "__main__":
    t1 = (torch.rand(1, 1, 32, 32, 32) * 32).int()
    print(t1.sum())
    out = decimal_to_bits(t1)

    dec = bits_to_decimal(out)

    print(dec.sum())
    print(out.shape)

    print(dec.shape)
